import torch


def gather(values: torch.Tensor, indices: torch.Tensor, dim: int):
    """
    :param values: shape=[..., n, ...]
    :param indices: shape=[..., m]
    :param dim: int
    :return: shape: shape=[..., m, ...]
    """
    dim %= values.ndim
    indices = torch.reshape(indices, indices.shape[:dim + 1] + (1,) * (values.ndim - dim - 1))
    indices = torch.broadcast_to(indices, indices.shape[:dim + 1] + values.shape[dim + 1:])
    return torch.gather(values, dim=dim, index=indices)


def scatter_set(values0: torch.Tensor, values: torch.Tensor, indices: torch.Tensor, dim: int):
    """
    :param values0: shape=[..., n, ...]
    :param values: shape=[..., m, ...]
    :param indices: shape=[..., m]
    :param dim: int
    :return: shape: shape=[..., n, ...]
    """
    dim %= values.ndim
    indices = torch.reshape(indices, indices.shape[:dim + 1] + (1,) * (values.ndim - dim - 1))
    indices = torch.broadcast_to(indices, indices.shape[:dim + 1] + values.shape[dim + 1:])
    return torch.scatter(values0, dim=dim, index=indices, src=values)


def scatter_add(values0: torch.Tensor, values: torch.Tensor, indices: torch.Tensor, dim: int):
    """
    :param values0: shape=[..., n, ...]
    :param values: shape=[..., m, ...]
    :param indices: shape=[..., m]
    :param dim: int
    :return: shape: shape=[..., n, ...]
    """
    dim %= values.ndim
    indices = torch.reshape(indices, indices.shape[:dim + 1] + (1,) * (values.ndim - dim - 1))
    indices = torch.broadcast_to(indices, indices.shape[:dim + 1] + values.shape[dim + 1:])
    return torch.scatter_add(values0, dim=dim, index=indices, src=values)
